import torch.nn as nn
from .config import *

class CNN18(nn.Module):
    def __init__(self, num_classes = idAmount):
        super(CNN18, self).__init__()

        # 为较小的灰度图像定制网络
        self.layers = nn.Sequential(
            # 初始卷积层调整为适应单通道输入，使用较小的卷积核
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),

            # 第一组卷积层，适度使用卷积核和步长
            *self._make_layer(32, 64, 2, stride=1),  # 较少的层数以适应18层总限制

            # 第二组卷积层
            *self._make_layer(64, 128, 2, stride=2),

            # 第三组卷积层
            *self._make_layer(128, 256, 2, stride=2),

            # 使用自适应平均池化完成空间尺寸压缩
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(256, num_classes)
        )

    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        layers = []
        # 跨层变化通道数，调整stride
        layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False))
        layers.append(nn.BatchNorm2d(out_channels))
        layers.append(nn.ReLU(inplace=True))

        # 添加更多的卷积层
        for _ in range(1, num_blocks):
            layers.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False))
            layers.append(nn.BatchNorm2d(out_channels))
            layers.append(nn.ReLU(inplace=True))

        return layers

    def forward(self, x):
        x = self.layers(x)
        return x


class CNN34(nn.Module):
    def __init__(self, num_classes=idAmount):
        super(CNN34, self).__init__()
        
        # 为较小的灰度图像定制网络
        self.layers = nn.Sequential(
            # 初始卷积层调整为适应单通道输入，使用较小的卷积核
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            
            # 第一组卷积层，适度使用卷积核和步长
            *self._make_layer(32, 64, 4, stride=1),
            
            # 第二组卷积层
            *self._make_layer(64, 128, 4, stride=2),
            
            # 第三组卷积层
            *self._make_layer(128, 256, 4, stride=2),
            
            # 避免过多的最大池化，使用自适应平均池化完成空间尺寸压缩
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(256, num_classes)
        )
    
    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        layers = []
        # 跨层变化通道数，调整stride
        layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False))
        layers.append(nn.BatchNorm2d(out_channels))
        layers.append(nn.ReLU(inplace=True))
        
        # 添加更多的卷积层
        for _ in range(1, num_blocks):
            layers.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False))
            layers.append(nn.BatchNorm2d(out_channels))
            layers.append(nn.ReLU(inplace=True))
        
        return layers

    def forward(self, x):
        x = self.layers(x)
        return x
    

class CNN50(nn.Module):
    def __init__(self, num_classes=idAmount):
        super(CNN50, self).__init__()
        
        # 为较小的灰度图像定制网络
        self.layers = nn.Sequential(
            # 初始卷积层调整为适应单通道输入，使用较小的卷积核
            nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            
            # 第一组卷积层
            *self._make_layer(32, 64, 8, stride=1),
            
            # 第二组卷积层
            *self._make_layer(64, 128, 12, stride=2),
            
            # 第三组卷积层
            *self._make_layer(128, 256, 16, stride=2),
            
            # 第四组卷积层
            *self._make_layer(256, 512, 12, stride=2),

            # 最终池化和全连接层
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(512, num_classes)
        )
    
    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
        layers = []
        # 首层在这一组内进行特征通道扩展和下采样
        layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False))
        layers.append(nn.BatchNorm2d(out_channels))
        layers.append(nn.ReLU(inplace=True))
        
        # 添加更多的卷积层以增加深度
        for _ in range(1, num_blocks):
            layers.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False))
            layers.append(nn.BatchNorm2d(out_channels))
            layers.append(nn.ReLU(inplace=True))
        
        return layers

    def forward(self, x):
        x = self.layers(x)
        return x
